/* -*- Mode:C++; c-file-style:"gnu"; indent-tabs-mode:nil; -*- */
/*
 * Copyright (c) 2013 Magister Solutions Ltd.
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License version 2 as
 * published by the Free Software Foundation;
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
 *
 * Author: Jani Puttonen <jani.puttonen@magister.fi>
 */

#include "satellite-return-link-encapsulator-arq.h"

#include "satellite-arq-buffer-context.h"
#include "satellite-arq-header.h"
#include "satellite-encap-pdu-status-tag.h"
#include "satellite-llc.h"
#include "satellite-mac-tag.h"
#include "satellite-queue.h"
#include "satellite-time-tag.h"

#include <ns3/log.h>
#include <ns3/mac48-address.h>
#include <ns3/simulator.h>

#include <algorithm>
#include <limits>
#include <map>
#include <utility>

NS_LOG_COMPONENT_DEFINE("SatReturnLinkEncapsulatorArq");

namespace ns3
{

NS_OBJECT_ENSURE_REGISTERED(SatReturnLinkEncapsulatorArq);

SatReturnLinkEncapsulatorArq::SatReturnLinkEncapsulatorArq()
    : m_seqNo(),
      m_txedBuffer(),
      m_retxBuffer(),
      m_retxBufferSize(0),
      m_txedBufferSize(0),
      m_maxRtnArqSegmentSize(37),
      m_maxNoOfRetransmissions(2),
      m_retransmissionTimer(Seconds(0.6)),
      m_arqWindowSize(10),
      m_arqHeaderSize(1),
      m_nextExpectedSeqNo(0)
{
    NS_LOG_FUNCTION(this);
    NS_ASSERT(false);

    /**
     * Default constructor is not meant to be used!
     */
}

SatReturnLinkEncapsulatorArq::SatReturnLinkEncapsulatorArq(Mac48Address encapAddress,
                                                           Mac48Address decapAddress,
                                                           Mac48Address sourceE2EAddress,
                                                           Mac48Address destE2EAddress,
                                                           uint8_t flowId,
                                                           uint32_t additionalHeaderSize)
    : SatReturnLinkEncapsulator(encapAddress,
                                decapAddress,
                                sourceE2EAddress,
                                destE2EAddress,
                                flowId,
                                additionalHeaderSize),
      m_seqNo(),
      m_txedBuffer(),
      m_retxBuffer(),
      m_retxBufferSize(0),
      m_txedBufferSize(0),
      m_maxRtnArqSegmentSize(37),
      m_maxNoOfRetransmissions(2),
      m_retransmissionTimer(Seconds(0.6)),
      m_arqWindowSize(10),
      m_arqHeaderSize(1),
      m_nextExpectedSeqNo(0)
{
    NS_LOG_FUNCTION(this);

    ObjectBase::ConstructSelf(AttributeConstructionList());

    m_seqNo = Create<SatArqSequenceNumber>(m_arqWindowSize);
}

SatReturnLinkEncapsulatorArq::~SatReturnLinkEncapsulatorArq()
{
    NS_LOG_FUNCTION(this);
}

TypeId
SatReturnLinkEncapsulatorArq::GetTypeId(void)
{
    static TypeId tid =
        TypeId("ns3::SatReturnLinkEncapsulatorArq")
            .SetParent<SatReturnLinkEncapsulator>()
            .AddConstructor<SatReturnLinkEncapsulatorArq>()
            .AddAttribute(
                "MaxRtnArqSegmentSize",
                "Maximum return link ARQ segment size in Bytes.",
                UintegerValue(37),
                MakeUintegerAccessor(&SatReturnLinkEncapsulatorArq::m_maxRtnArqSegmentSize),
                MakeUintegerChecker<uint32_t>())
            .AddAttribute(
                "MaxNoOfRetransmissions",
                "Maximum number of retransmissions for a single RLE PDU.",
                UintegerValue(2),
                MakeUintegerAccessor(&SatReturnLinkEncapsulatorArq::m_maxNoOfRetransmissions),
                MakeUintegerChecker<uint32_t>())
            .AddAttribute(
                "RetransmissionTimer",
                "Retransmission time value, i.e. how long to wait for ACK before retransmission.",
                TimeValue(Seconds(0.6)),
                MakeTimeAccessor(&SatReturnLinkEncapsulatorArq::m_retransmissionTimer),
                MakeTimeChecker())
            .AddAttribute(
                "WindowSize",
                "Window size for ARQ, i.e. how many simultaneous packets are allowed in the air.",
                UintegerValue(10),
                MakeUintegerAccessor(&SatReturnLinkEncapsulatorArq::m_arqWindowSize),
                MakeUintegerChecker<uint32_t>())
            .AddAttribute("ArqHeaderSize",
                          "ARQ header size in Bytes.",
                          UintegerValue(1),
                          MakeUintegerAccessor(&SatReturnLinkEncapsulatorArq::m_arqHeaderSize),
                          MakeUintegerChecker<uint32_t>())
            .AddAttribute("RxWaitingTime",
                          "Time to wait for a packet at the reception (GW) before moving onwards "
                          "with the packet reception.",
                          TimeValue(Seconds(1.8)),
                          MakeTimeAccessor(&SatReturnLinkEncapsulatorArq::m_rxWaitingTimer),
                          MakeTimeChecker());
    return tid;
}

TypeId
SatReturnLinkEncapsulatorArq::GetInstanceTypeId() const
{
    return GetTypeId();
}

void
SatReturnLinkEncapsulatorArq::DoDispose()
{
    NS_LOG_FUNCTION(this);
    m_seqNo = 0;

    // Clean-up the Tx'ed buffer
    std::map<uint8_t, Ptr<SatArqBufferContext>>::iterator it = m_txedBuffer.begin();
    while (it != m_txedBuffer.end())
    {
        it->second->DoDispose();
        it->second = 0;
        ++it;
    }
    m_txedBuffer.clear();

    // Clean-up the reTx buffer
    it = m_retxBuffer.begin();
    while (it != m_retxBuffer.end())
    {
        it->second->DoDispose();
        it->second = 0;
        ++it;
    }
    m_retxBuffer.clear();

    // Clean-up the reordering buffer
    std::map<uint32_t, Ptr<SatArqBufferContext>>::iterator it2 = m_reorderingBuffer.begin();
    while (it2 != m_reorderingBuffer.end())
    {
        it2->second->DoDispose();
        it2->second = 0;
        ++it2;
    }
    m_reorderingBuffer.clear();

    SatReturnLinkEncapsulator::DoDispose();
}

Ptr<Packet>
SatReturnLinkEncapsulatorArq::NotifyTxOpportunity(uint32_t bytes,
                                                  uint32_t& bytesLeft,
                                                  uint32_t& nextMinTxO)
{
    NS_LOG_FUNCTION(this << bytes);
    NS_LOG_INFO("TxOpportunity for UT: " << m_encapAddress << " flowId: " << (uint32_t)m_flowId
                                         << " of " << bytes << " bytes");

    // Payload adapted PDU = NULL
    Ptr<Packet> packet;

    /**
     * Give priority to the retransmissions. Whenever, the retransmission
     * timer is expired, packet is moved to the retransmission buffer from
     * the transmitted buffer.
     */
    if (!m_retxBuffer.empty())
    {
        // Oldest seqNo sent first
        Ptr<SatArqBufferContext> context = m_retxBuffer.begin()->second;

        // If the packet fits into the transmission opportunity
        if (context->m_pdu->GetSize() <= bytes)
        {
            // Pop the front
            m_retxBuffer.erase(m_retxBuffer.begin());

            // Increase the retransmission counter
            context->m_retransmissionCount = context->m_retransmissionCount + 1;

            m_retxBufferSize -= context->m_pdu->GetSize();
            m_txedBufferSize += context->m_pdu->GetSize();

            // Store it back to the transmitted packet container.
            m_txedBuffer.insert(std::make_pair(context->m_seqNo, context));

            // Create the retransmission event and store it to the context. Event is cancelled if a
            // ACK is received. However, if the event triggers, we shall send the packet again, if
            // the packet still has retransmissions left.
            EventId t = Simulator::Schedule(m_retransmissionTimer,
                                            &SatReturnLinkEncapsulatorArq::ArqReTxTimerExpired,
                                            this,
                                            context->m_seqNo);
            context->m_waitingTimer = t;

            NS_LOG_INFO("UT: << " << m_encapAddress << " sent a retransmission packet of size: "
                                  << context->m_pdu->GetSize()
                                  << " with seqNo: " << (uint32_t)(context->m_seqNo)
                                  << " flowId: " << (uint32_t)(m_flowId));

            Ptr<Packet> copy = context->m_pdu->Copy();
            return copy;
        }
        else
        {
            NS_LOG_INFO("Retransmission PDU: " << context->m_pdu->GetUid()
                                               << " size: " << context->m_pdu->GetSize()
                                               << " does not fit into TxO: " << bytes);
        }
    }

    // Check the transmission buffer. Sequence number needs to be
    // available for any new transmissions.
    else if (!m_txQueue->IsEmpty() && m_seqNo->SeqNoAvailable())
    {
        // Crate new RLE PDU
        packet = GetNewRlePdu(bytes, m_maxRtnArqSegmentSize, m_arqHeaderSize);

        if (packet)
        {
            // Add MAC tag to identify the packet in lower layers
            SatMacTag mTag;
            if (!packet->PeekPacketTag(mTag))
            {
                mTag.SetDestAddress(m_decapAddress);
                mTag.SetSourceAddress(m_encapAddress);
                packet->AddPacketTag(mTag);
            }

            // Add E2E address tag to identify the packet in lower layers
            SatAddressE2ETag addressE2ETag;
            if (!packet->PeekPacketTag(addressE2ETag))
            {
                addressE2ETag.SetE2EDestAddress(m_destE2EAddress);
                addressE2ETag.SetE2ESourceAddress(m_sourceE2EAddress);
                packet->AddPacketTag(addressE2ETag);
            }

            // Add flow id tag
            SatFlowIdTag flowIdTag;
            flowIdTag.SetFlowId(m_flowId);
            packet->AddPacketTag(flowIdTag);

            // Get next available sequence number
            uint8_t seqNo = m_seqNo->NextSequenceNumber();

            // Add ARQ header
            SatArqHeader arqHeader;
            arqHeader.SetSeqNo(seqNo);
            packet->AddHeader(arqHeader);

            // Create ARQ context and store it to Tx'ed buffer
            Ptr<SatArqBufferContext> arqContext = CreateObject<SatArqBufferContext>();
            arqContext->m_retransmissionCount = 0;
            Ptr<Packet> copy = packet->Copy();
            arqContext->m_pdu = copy;
            arqContext->m_seqNo = seqNo;

            // Create the retransmission event and store it to the context. Event is cancelled if a
            // ACK is received. However, if the event triggers, we shall send the packet again, if
            // the packet still has retransmissions left.
            arqContext->m_waitingTimer =
                Simulator::Schedule(m_retransmissionTimer,
                                    &SatReturnLinkEncapsulatorArq::ArqReTxTimerExpired,
                                    this,
                                    seqNo);

            // Update the buffer status
            m_txedBufferSize += packet->GetSize();
            m_txedBuffer.insert(std::make_pair(seqNo, arqContext));

            if (packet->GetSize() > bytes)
            {
                NS_FATAL_ERROR("Created packet of size: " << packet->GetSize()
                                                          << " is larger than the tx opportunity: "
                                                          << bytes);
            }

            NS_LOG_INFO("UT: << " << m_encapAddress << " sent a packet of size: "
                                  << packet->GetSize() << " with seqNo: " << (uint32_t)(seqNo)
                                  << " flowId: " << (uint32_t)(m_flowId));
            NS_LOG_INFO("Queue size after TxOpportunity: " << m_txQueue->GetNBytes());
        }
    }
    else if (!m_seqNo->SeqNoAvailable())
    {
        bytesLeft = 0;
        return packet;
    }

    // Update bytes lefts
    bytesLeft = GetTxBufferSizeInBytes();

    // Update min TxO
    nextMinTxO = GetMinTxOpportunityInBytes();

    return packet;
}

void
SatReturnLinkEncapsulatorArq::ArqReTxTimerExpired(uint8_t seqNo)
{
    NS_LOG_FUNCTION(this << (uint32_t)seqNo);

    NS_LOG_INFO("At UT: " << m_encapAddress
                          << " ARQ retransmission timer expired for: " << (uint32_t)(seqNo));

    std::map<uint8_t, Ptr<SatArqBufferContext>>::iterator it = m_txedBuffer.find(seqNo);

    if (it != m_txedBuffer.end())
    {
        NS_ASSERT(seqNo == it->second->m_seqNo);
        NS_ASSERT(it->second->m_pdu);

        // Retransmission still possible
        if (it->second->m_retransmissionCount < m_maxNoOfRetransmissions)
        {
            NS_LOG_INFO("Moving the ARQ context to retransmission buffer");

            Ptr<SatArqBufferContext> context = it->second;

            m_txedBuffer.erase(it);
            m_retxBufferSize += context->m_pdu->GetSize();

            // Push to the retransmission buffer
            m_retxBuffer.insert(std::make_pair(seqNo, context));
        }
        // Maximum retransmissions reached
        else
        {
            NS_LOG_INFO("For UT: " << m_encapAddress << " max retransmissions reached for "
                                   << (uint32_t)(seqNo));

            // Do clean-up
            CleanUp(seqNo);
        }
    }
    else
    {
        NS_LOG_INFO("Element not found anymore in the m_txedBuffer, thus ACK has been received "
                    "already earlier");
    }
}

void
SatReturnLinkEncapsulatorArq::CleanUp(uint8_t sequenceNumber)
{
    NS_LOG_FUNCTION(this << (uint32_t)sequenceNumber);

    // Release sequence number
    m_seqNo->Release(sequenceNumber);

    // Clean-up the Tx'ed buffer
    std::map<uint8_t, Ptr<SatArqBufferContext>>::iterator it = m_txedBuffer.find(sequenceNumber);
    if (it != m_txedBuffer.end())
    {
        NS_LOG_INFO("Sequence no: " << (uint32_t)sequenceNumber << " clean up from txedBuffer!");
        m_txedBufferSize -= it->second->m_pdu->GetSize();
        it->second->DoDispose();
        it->second = 0;
        m_txedBuffer.erase(it);
    }

    // Clean-up the reTx buffer
    it = m_retxBuffer.find(sequenceNumber);
    if (it != m_retxBuffer.end())
    {
        NS_LOG_INFO("Sequence no: " << (uint32_t)sequenceNumber << " clean up from retxBuffer!");
        m_retxBufferSize -= it->second->m_pdu->GetSize();
        it->second->DoDispose();
        it->second = 0;
        m_retxBuffer.erase(it);
    }
}

void
SatReturnLinkEncapsulatorArq::ReceiveAck(Ptr<SatArqAckMessage> ack)
{
    NS_LOG_FUNCTION(this);

    /**
     * Take the packet away from the transmitted or retransmission
     * buffer and release the sequence number for further use.
     */

    NS_LOG_INFO("UT: " << m_encapAddress
                       << " received ACK with SN: " << (uint32_t)(ack->GetSequenceNumber()));

    // Do clean-up
    CleanUp(ack->GetSequenceNumber());
}

void
SatReturnLinkEncapsulatorArq::ReceivePdu(Ptr<Packet> p)
{
    NS_LOG_FUNCTION(this << p->GetSize());

    // Remove encap PDU status tag
    SatEncapPduStatusTag statusTag;
    p->RemovePacketTag(statusTag);

    // Remove flow id tag
    SatFlowIdTag flowIdTag;
    p->RemovePacketTag(flowIdTag);

    // Sanity check
    SatMacTag mTag;
    bool mSuccess = p->RemovePacketTag(mTag);
    if (!mSuccess)
    {
        NS_FATAL_ERROR("MAC tag not found in the packet!");
    }
    else if (mTag.GetDestAddress() != m_decapAddress)
    {
        NS_FATAL_ERROR("Packet was not intended for this receiver!");
    }

    SatArqHeader arqHeader;
    p->RemoveHeader(arqHeader);
    uint8_t seqNo = arqHeader.GetSeqNo();

    NS_LOG_INFO("UT: " << m_encapAddress << " received a packet with SeqNo: " << (uint32_t)(seqNo));

    // Send ACK for the received RLE packet.
    SendAck(seqNo);

    // Convert the 8-bit sequence number to continuous 32-bit sequence number
    uint32_t sn = ConvertSeqNo(seqNo);

    NS_LOG_INFO("8bit SN: " << (uint32_t)(seqNo) << " 32bit SN: " << sn);

    // If the received SN is valid. If we receive a SN from the past
    // nothing is needed to be done.
    if (sn >= m_nextExpectedSeqNo)
    {
        std::map<uint32_t, Ptr<SatArqBufferContext>>::iterator it = m_reorderingBuffer.find(sn);

        // If the context is not found, then we create a new one.
        if (it == m_reorderingBuffer.end())
        {
            NS_LOG_INFO("UT: " << m_encapAddress
                               << " created a new ARQ buffer entry for SeqNo: " << sn);
            Ptr<SatArqBufferContext> arqContext = CreateObject<SatArqBufferContext>();
            arqContext->m_pdu = p;
            arqContext->m_rxStatus = true;
            arqContext->m_seqNo = sn;
            arqContext->m_retransmissionCount = 0;
            m_reorderingBuffer.insert(std::make_pair(sn, arqContext));
        }
        // If the context is found, update it.
        else
        {
            NS_LOG_INFO("UT: " << m_encapAddress
                               << " reset an existing ARQ entry for SeqNo: " << sn);
            it->second->m_waitingTimer.Cancel();
            it->second->m_pdu = p;
            it->second->m_rxStatus = true;
        }

        NS_LOG_INFO("Received a packet with SeqNo: " << sn
                                                     << ", expecting: " << m_nextExpectedSeqNo);

        // If this is not the SN we expect
        if (sn != m_nextExpectedSeqNo)
        {
            // Add context
            for (uint32_t i = m_nextExpectedSeqNo; i < sn; ++i)
            {
                std::map<uint32_t, Ptr<SatArqBufferContext>>::iterator it2 =
                    m_reorderingBuffer.find(i);

                NS_LOG_INFO("Finding context for " << i);

                // If context not found
                if (it2 == m_reorderingBuffer.end())
                {
                    NS_LOG_INFO("Context NOT found for SeqNo: " << i);

                    Ptr<SatArqBufferContext> arqContext = CreateObject<SatArqBufferContext>();
                    arqContext->m_pdu = nullptr;
                    arqContext->m_rxStatus = false;
                    arqContext->m_seqNo = i;
                    arqContext->m_retransmissionCount = 0;
                    m_reorderingBuffer.insert(std::make_pair(i, arqContext));
                    EventId id =
                        Simulator::Schedule(m_rxWaitingTimer,
                                            &SatReturnLinkEncapsulatorArq::RxWaitingTimerExpired,
                                            this,
                                            i);
                    arqContext->m_waitingTimer = id;
                }
            }
        }
        // An expected sequence number received, reassemble and receive.
        else
        {
            ReassembleAndReceive();
        }
    }
    else
    {
        NS_LOG_INFO("UT: " << m_encapAddress << " received a packet with SeqNo: " << sn
                           << " which is already received!");
    }
}

uint32_t
SatReturnLinkEncapsulatorArq::ConvertSeqNo(uint8_t seqNo) const
{
    NS_LOG_FUNCTION(this << (uint32_t)seqNo);

    uint32_t globalSeqNo(0);

    // Calculate the rounds and current seq no from m_nextExpectedSeqNo
    uint32_t rounds = (m_nextExpectedSeqNo / std::numeric_limits<uint8_t>::max());
    uint32_t rawSeqNo = m_nextExpectedSeqNo % std::numeric_limits<uint8_t>::max();

    NS_LOG_INFO("Input: " << (uint32_t)(seqNo) << " rounds: " << rounds << " rawSeqNo: " << rawSeqNo
                          << " windowSize: " << m_arqWindowSize
                          << " next expected: " << m_nextExpectedSeqNo);

    // Received sequence number is higher than the expected one.
    if (seqNo >= rawSeqNo)
    {
        // If seqNo is from previous round
        if ((seqNo - rawSeqNo) > 2 * m_arqWindowSize)
        {
            rounds--;
        }
    }
    // seqNo < rawSeqNo
    else
    {
        if ((rawSeqNo - seqNo) > 2 * m_arqWindowSize)
        {
            rounds++;
        }
    }

    globalSeqNo = rounds * std::numeric_limits<uint8_t>::max() + seqNo;

    return globalSeqNo;
}

void
SatReturnLinkEncapsulatorArq::ReassembleAndReceive()
{
    NS_LOG_FUNCTION(this);

    // Start from the expected sequence number iterator
    std::map<uint32_t, Ptr<SatArqBufferContext>>::iterator it = m_reorderingBuffer.begin();

    NS_LOG_INFO("Process SeqNo: " << it->first << ", expected: " << m_nextExpectedSeqNo
                                  << ", status: " << it->second->m_rxStatus);

    /**
     * As long as the PDU is the next expected one, process the PDU
     * and erase it.
     */
    while (it != m_reorderingBuffer.end() && it->first == m_nextExpectedSeqNo &&
           it->second->m_rxStatus == true)
    {
        NS_LOG_INFO("Process SeqNo: " << it->first << ", expected: " << m_nextExpectedSeqNo
                                      << ", status: " << it->second->m_rxStatus);

        // If timer is running, cancel it.
        if (it->second->m_waitingTimer.IsPending())
        {
            it->second->m_waitingTimer.Cancel();
        }

        // If PDU == NULL, it means that the RxWaitingTimer has expired
        // without PDU being received
        if (it->second->m_pdu)
        {
            // Process the PDU
            ProcessPdu(it->second->m_pdu);
        }

        m_reorderingBuffer.erase(it);
        it = m_reorderingBuffer.begin();

        // Increase the seq no
        ++m_nextExpectedSeqNo;

        NS_LOG_INFO("Increasing SeqNo to " << m_nextExpectedSeqNo);
    }
}

void
SatReturnLinkEncapsulatorArq::RxWaitingTimerExpired(uint32_t seqNo)
{
    NS_LOG_FUNCTION(this << seqNo);

    NS_LOG_INFO("For UT: " << m_encapAddress << " max waiting time reached for SeqNo: " << seqNo);
    NS_LOG_INFO("Mark the PDU received and move forward!");

    // Find waiting timer, erase it and mark the packet received.
    std::map<uint32_t, Ptr<SatArqBufferContext>>::iterator it = m_reorderingBuffer.find(seqNo);
    if (it != m_reorderingBuffer.end())
    {
        it->second->m_waitingTimer.Cancel();
        it->second->m_rxStatus = true;
    }
    else
    {
        NS_FATAL_ERROR("Rx waiting timer is not running anymore even though it expired!");
    }

    ReassembleAndReceive();
}

uint32_t
SatReturnLinkEncapsulatorArq::GetTxBufferSizeInBytes() const
{
    NS_LOG_FUNCTION(this);

    return m_txQueue->GetNBytes() + m_retxBufferSize;
}

void
SatReturnLinkEncapsulatorArq::SendAck(uint8_t seqNo) const
{
    NS_LOG_FUNCTION(this << (uint32_t)seqNo);

    NS_LOG_INFO("GW: " << m_decapAddress << " send ACK to UT: " << m_encapAddress
                       << " with flowId: " << (uint32_t)(m_flowId)
                       << " with SN: " << (uint32_t)(seqNo));

    /**
     * RLE sends the ACK control message via a callback to SatNetDevice of the GW to the
     * GSE control buffer. The transmission assumptions for the control messages are currently
     * - Default (most robust) MODCOD
     * - Strict priority
     */
    if (!m_ctrlCallback.IsNull())
    {
        Ptr<SatArqAckMessage> ack = Create<SatArqAckMessage>();
        ack->SetSequenceNumber(seqNo);
        ack->SetFlowId(m_flowId);

        // Source address (UT) is used here, since the in RTN the GW is
        // sending the ACK to the UT.
        m_ctrlCallback(ack, m_encapAddress);
    }
    else
    {
        NS_FATAL_ERROR("Unable to send ACK, since the Ctrl callback is NULL!");
    }
}

} // namespace ns3
